import tensorflow as tf
from tensorflow.keras import layers, models
import pandas as pd
from sklearn.preprocessing import MultiLabelBinarizer
from sklearn.model_selection import train_test_split
from sklearn.metrics import roc_auc_score

# 自定义TensorBoard AUC记录回调
class AUCLoggingCallback(tf.keras.callbacks.Callback):
    def __init__(self, log_dir, x_val, y_val):
        super(AUCLoggingCallback, self).__init__()
        self.log_dir = log_dir
        self.x_val = x_val
        self.y_val = y_val
        self.writer = tf.summary.create_file_writer(log_dir)

    def on_epoch_end(self, epoch, logs=None):
        y_pred = self.model.predict(self.x_val)
        # 转换为二元分类
        y_pred_binary = (y_pred > 3.5).astype(int)
        y_true_binary = (self.y_val > 3.5).astype(int)
        # 计算AUC
        auc_value = roc_auc_score(y_true_binary, y_pred_binary)
        with self.writer.as_default():
            tf.summary.scalar('auc', auc_value, step=epoch)

# 加载数据集
movies_df = pd.read_csv('movies.csv')
ratings_df = pd.read_csv('ratings.csv')

# 处理电影类型数据
genres = movies_df['genres'].apply(lambda x: x.split('|'))
mlb = MultiLabelBinarizer()
genres_encoded = mlb.fit_transform(genres)
movies_df = movies_df.join(pd.DataFrame(genres_encoded, index=movies_df.index))

# 合并电影和评分数据
data = ratings_df.merge(movies_df, on='movieId')

# 提取特征和标签
X = data.drop(['userId', 'movieId', 'rating', 'timestamp', 'title', 'genres'], axis=1)
y = data['rating']

# 分割数据集
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# 创建电影的特征输入
movie_input = tf.keras.Input(shape=(X.shape[1],), name='Movie-Input')

# 创建网络层
fc1 = layers.Dense(units=128, activation='relu')(movie_input)
dropout = layers.Dropout(0.5)(fc1)
fc2 = layers.Dense(units=64, activation='relu')(dropout)
fc3 = layers.Dense(units=32, activation='relu')(fc2)

# 输出层
output = layers.Dense(units=1)(fc3)  # 评分作为回归任务的输出

# 创建模型
model = models.Model(inputs=movie_input, outputs=output)
model.compile(optimizer='adam', loss='mean_squared_error', metrics=['mae'])
log_dir='./log/seq'
# 训练模型
auc_logging_callback = AUCLoggingCallback(log_dir, X_test, y_test)
tensorboard_callback = tf.keras.callbacks.TensorBoard(log_dir, histogram_freq=1)
model.fit(X_train, y_train, validation_data=(X_test, y_test), epochs=10,
          callbacks=[tensorboard_callback,auc_logging_callback])




# 推荐函数
def recommend(user_id, top_n=5):
    # 获取用户评过分的电影
    user_rated_movies = ratings_df[ratings_df['userId'] == user_id]
    # 为用户未评分的电影创建预测集
    user_unrated_movies = movies_df[~movies_df['movieId'].isin(user_rated_movies['movieId'])]
    user_unrated_genres = user_unrated_movies.drop(['movieId', 'title', 'genres'], axis=1)
    # 预测评分
    predictions = model.predict(user_unrated_genres)
    user_unrated_movies['predicted_rating'] = predictions
    # 排序并返回预测评分最高的电影
    recommended_movies = user_unrated_movies.sort_values(by='predicted_rating', ascending=False).head(top_n)
    return recommended_movies[['title', 'predicted_rating']]

# 为userId为1的用户推荐电影
print(recommend(1))
